package org.dromara.common.core.utils;

import cn.hutool.core.collection.CollectionUtil;
import lombok.Data;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @author mokj3
 * @date 2021/11/19 14:43
 * @description 树结构构建工具
 */
public class TreeUtils {

    /**
     * @param source   元数据集合
     * @param mFunc    模型转换函数
     * @param keyFunc  key的获取函数
     * @param pKeyFunc 父级key的获取函数
     * @param <K>      key的类型
     * @param <T>      原目标类型
     * @param <R>      结果模型类型
     * @return 返回内部类，树结构
     */
    public static <K, T, R> List<TreeNode<R>> build(Collection<T> source, Function<? super T, R> mFunc, Function<? super R, K> keyFunc, Function<? super R, K> pKeyFunc) {
        //转树节点模型
        List<TreeNode<R>> collect = source.stream()
                .map(mode -> new TreeNode<>(mFunc.apply(mode)))
                .collect(Collectors.toList());
        //根据父级key分组
        Map<K, List<TreeNode<R>>> listMap = collect.stream()
                .collect(Collectors.groupingBy(node -> pKeyFunc.apply(node.curNode)));
        //当前节点与孩子节点建立联系
        collect.forEach(rTreeNode ->
                rTreeNode.children = listMap.get(keyFunc.apply(rTreeNode.curNode)));
        return collect;
    }


    /**
     * 把数据集转换成树结构数据
     *
     * @param source   数据源
     * @param m1Func   模型转换函数：源数据转树模型
     * @param keyFunc  唯一key
     * @param pKeyFunc 父级key
     * @param filter   过滤函数：过滤出你想要的树结构数据
     * @param m2Func   模型转换函数：树模型转数据结果模型
     * @param sFunc    当前节点与孩子节点，关系处理函数
     * @param <K>      key的类型：数字、字符串
     * @param <T>      源数据的类型
     * @param <R1>     内部树结构模型
     * @param <R2>     外部展示的树结构模型
     * @return 最终树结果集
     */
    public static <K, T, R1, R2> List<R2> buildTree(Collection<T> source, Function<? super T, R1> m1Func,
                                                    Function<? super R1, K> keyFunc, Function<? super R1, K> pKeyFunc,
                                                    Predicate<? super R1> filter,
                                                    Function<R1, R2> m2Func, BiConsumer<R2, List<R2>> sFunc) {
        //把源数据构建树结构，建立父子级关系
        List<TreeNode<R1>> treeNodes = build(source, m1Func, keyFunc, pKeyFunc);
        //过滤出想要的节点，例：过滤出所有根节点
        List<TreeNode<R1>> treeNodeList = treeNodes.stream()
                .filter(tTreeNode -> filter.test(tTreeNode.curNode))
                .collect(Collectors.toList());
        //内部模型转换成可供展示调用的模型结构
        return TreeUtils.convertNodes(treeNodeList, m2Func, sFunc);
    }


    public static <K, T> List<T> buildTree(Collection<T> source,
                                           Function<? super T, K> keyFunc, Function<? super T, K> pKeyFunc,
                                           Predicate<? super T> filter, BiConsumer<T, List<T>> sFunc) {
        return buildTree(source, Function.identity(), keyFunc, pKeyFunc, filter, Function.identity(), sFunc);
    }

    /**
     * 转换树结构模型，从内部类转换成可对外展示调用的模型
     *
     * @param treeNodes 节点集合
     * @param cFunc     内部类模型转变外部模型函数：就是告诉我-新模型和老模型的转换过程
     * @param sFunc     子节点存放位置外部模型函数：就是告诉我-新模型的孩子节点是哪个字段
     * @param <T>       老模型的类型
     * @param <R>       新模型的类型
     * @return 新模型结果集合
     */
    public static <T, R> List<R> convertNodes(List<TreeNode<T>> treeNodes, Function<T, R> cFunc, BiConsumer<R, List<R>> sFunc) {
        return treeNodes.stream().map(treeNode -> convertNode(treeNode, cFunc, sFunc)).collect(Collectors.toList());
    }

    /**
     * 转换树结构模型，从内部类转换成可对外展示调用的模型
     *
     * @param treeNode 节点
     * @param cFunc    内部类模型转变外部模型函数：就是告诉我-新模型和老模型的转换过程
     * @param sFunc    子节点存放位置外部模型函数：就是告诉我-新模型的孩子节点是哪个字段
     * @param <T>      老模型的类型
     * @param <R>      新模型的类型
     * @return 新模型结果
     */
    public static <T, R> R convertNode(TreeNode<T> treeNode, Function<T, R> cFunc, BiConsumer<R, List<R>> sFunc) {
        List<TreeNode<T>> children = treeNode.getChildren();
        ArrayList<R> arrayList = new ArrayList<>();
        if (CollectionUtil.isNotEmpty(children)) {
            for (TreeNode<T> child : children) {
                //递归去转换模型
                arrayList.add(convertNode(child, cFunc, sFunc));
            }
        }
        //类型转换
        R responseDTO = cFunc.apply(treeNode.curNode);
        //新模型与新生成的孩子节点，需关系处理
        sFunc.accept(responseDTO, arrayList);

        return responseDTO;
    }

    /**
     * 获取节点的子节点
     *
     * @param treeNode 一个节点
     * @param <T>      节点类型
     * @return 结果集
     */
    public static <T> List<T> getChildren(TreeNode<T> treeNode) {
        List<TreeNode<T>> children = treeNode.getChildren();
        if (CollectionUtil.isEmpty(children)) {
            return new ArrayList<>();
        }
        return children.stream().map(TreeNode::getCurNode).collect(Collectors.toList());
    }


    @Data
    public static class TreeNode<R> {
        //当前节点
        private R curNode;
        //子节点
        private List<TreeNode<R>> children;

        public TreeNode() {
        }

        public TreeNode(R curNode) {
            this.curNode = curNode;
        }
    }

}
